#include <string>
#include <stack>
#include <map>
#include <iterator>

using std::string;
using std::stack;
using std::map;
/*
 * formula = "K4(ON(SO3)2)2"
 * Output: "K4N2O14S4"
 * Explanation: 
 * The count of elements are {'K': 4, 'N': 2, 'O': 14, 'S': 4}.
 */

string flattenString(string s) {
    string result;
    int len = s.size();
    int level = 0;
    string sub_s = "";
    string numeral_s = "";
    stack<string> sub_stack;

    for (int i = 0; i < len; i++) {
        if ((s[i] >= 'A' && s[i] <= 'Z') ||
            (s[i] >= 'a' && s[i] <= 'z') ||
            (s[i] >= '0' && s[i] <= '9')) {
            if (level == 0) {
                if (sub_s != "") {
                    result += sub_s;
                    sub_s = "";
                }
                result += s[i];
            } else {
                sub_s += s[i];
            }
        }
        if (s[i] == '(') {
            sub_stack.push(sub_s);
            sub_s = "";
            level++;
        }
        if (s[i] == ')') {
            // get: xxx)<_out_number>
            int _out_number;
            string _out_number_str = string(1, s[i+1]);
            int ni = i + 2;
            for (; ni < len; ni++) {
                if (s[ni] >= '0' && s[ni] <= '9') {
                    _out_number_str.push_back(s[ni]);
                } else {
                    break;
                }
            }
            _out_number = std::stoi(_out_number_str);

            string tmp_s = sub_stack.top();
            // parse sub_string
            int l = sub_s.size();
            string _number_str = "";
            for (int ii = 1; ii < l; ii++) {
                if (sub_s[ii] >= 'A' && sub_s[ii] <= 'Z') {
                    if (_number_str != "") {
                        int _number = std::stoi(_number_str); // 12
                        tmp_s += std::to_string(_number * _out_number);
                        _number_str = "";
                    } else {
                        tmp_s.push_back(sub_s[ii-1]);
                        tmp_s += _out_number_str;
                    }
                } else if (sub_s[ii] >= 'a' && sub_s[ii] <= 'z') {
                    tmp_s.push_back(sub_s[ii-1]);
                } else {
                    if (_number_str == "") {
                        tmp_s.push_back(sub_s[ii-1]);
                    }
                    _number_str.push_back(sub_s[ii]);
                }
                if (ii == l - 1) {
                    if (_number_str != "") {
                        int _number = std::stoi(_number_str);
                        tmp_s += std::to_string(_number * _out_number);
                    } else {
                        tmp_s.push_back(sub_s[ii]);
                        tmp_s += _out_number_str;
                    }
                }
            }
            sub_s = tmp_s;
            level--;
            i += _out_number_str.size();
            sub_stack.pop();
        }
    }
    result += sub_s;
    return result;
}

string numberOfAtoms(string formula) {
    if (formula == "") {
        return "";
    }

    // "K4(ON(SO3)2)2" -> "K4O2N2S4O12"
    string flatten_s = flattenString(formula);

    map<string, int> m_atom_number;
    string result;

    int len = flatten_s.size();
    string _number_str = "";
    string _atom_str = string(1, flatten_s[0]);
    for (int i = 1; i < len; i++) {
        if (flatten_s[i] >= 'A' && flatten_s[i] <= 'Z') {
            // O12K at K; _number_str: 12
            m_atom_number[_atom_str] += (
                    _number_str == "" ? 1 : std::stoi(_number_str));
            _number_str = "";
            _atom_str = string(1, flatten_s[i]);
        }
        if (flatten_s[i] >= 'a' && flatten_s[i] <= 'z') {
            _atom_str.push_back(flatten_s[i]);
        }
        if (flatten_s[i] >= '0' && flatten_s[i] <= '9') {
            _number_str.push_back(flatten_s[i]);
        }
    }
    if (_number_str != "") {
        m_atom_number[_atom_str] += std::stoi(_number_str);
    } else {
        m_atom_number[_atom_str] += 1;
    }

    map<string, int>::iterator itr;
    for (itr = m_atom_number.begin(); itr != m_atom_number.end(); ++itr) {
        result += itr->first;
        if (itr->second > 1) {
            result += std::to_string(itr->second);
        }
    }

    return result;
}

int main() {
    string formula = "Mg(OH)2K4(ON(AlSO3)2)2";
    string out = numberOfAtoms(formula);
    printf("%s -> %s\n", formula.c_str(), out.c_str());
    formula = "H50";
    out = numberOfAtoms(formula);
    printf("%s -> %s\n", formula.c_str(), out.c_str());
    formula = "N";
    out = numberOfAtoms(formula);
    printf("%s -> %s\n", formula.c_str(), out.c_str());
}
